import numpy as np
import os
import struct

def gen_golden_data_constant_pad():
    # 生成输入数据 [8, 1, 9, 2045] - 对应填充前的尺寸
    input_x = np.random.uniform(-2, 2, [8, 1, 9, 2045]).astype(np.float32)
    
    # 填充参数
    pad_left = 1
    pad_right = 2  
    pad_top = 3
    pad_bottom = 4
    pad_value = 0.0
    
    # 使用numpy进行填充
    pad_width = ((0, 0), (0, 0), (pad_top, pad_bottom), (pad_left, pad_right))
    golden = np.pad(input_x, pad_width, mode='constant', constant_values=pad_value)
    
    print(f"输入形状: {input_x.shape}")
    print(f"输出形状: {golden.shape}")
    print(f"填充参数: 左{pad_left}, 右{pad_right}, 上{pad_top}, 下{pad_bottom}, 值{pad_value}")
    
    # 创建输入输出目录
    os.makedirs("./input", exist_ok=True)
    os.makedirs("./output", exist_ok=True)
    
    # 根据 ConstantPadCustomTilingData 结构体生成tiling数据
    # struct ConstantPadCustomTilingData {
    #     uint32_t totalRows;  // outputN * outputC * outputH
    #     uint32_t tileNum;
    # 
    #     // 形状设置
    #     int32_t outputN;
    #     int32_t outputC;
    #     int32_t outputH;
    #     int32_t outputW;
    # 
    #     // 填充参数
    #     int32_t pad_left;
    #     int32_t pad_right;
    #     int32_t pad_top;
    #     int32_t pad_bottom;
    #     float pad_value;
    # };
    
    # 计算输出形状
    outputN, outputC, outputH, outputW = golden.shape
    total_rows = outputN * outputC * outputH  # totalRows
    tile_num = 8  # 8个tile对应8个核
    
    # 使用struct.pack确保正确的内存布局
    # 'I' = unsigned int (uint32_t), 'i' = int (int32_t), 'f' = float
    tiling_data = struct.pack('IIiiiiiiiif', 
                             total_rows,      # uint32_t totalRows
                             tile_num,        # uint32_t tileNum
                             outputN,         # int32_t outputN
                             outputC,         # int32_t outputC  
                             outputH,         # int32_t outputH
                             outputW,         # int32_t outputW
                             pad_left,        # int32_t pad_left
                             pad_right,       # int32_t pad_right
                             pad_top,         # int32_t pad_top
                             pad_bottom,      # int32_t pad_bottom
                             pad_value)       # float pad_value
    
    # 保存tiling数据
    with open("./input/input_tiling.bin", "wb") as f:
        f.write(tiling_data)
    
    # 保存输入和输出数据
    input_x.tofile("./input/input_x.bin")
    golden.tofile("./output/golden.bin")

if __name__ == "__main__":
    gen_golden_data_constant_pad()